import json
import os

import click
import dotenv
import openai
import pandas as pd
from tqdm import tqdm
from tree_rag.agents import DEFAULT_MODEL_NAME
from tree_rag.client import get_by_name
from tree_rag.dataclasses.knowledge import Scene

dotenv.load_dotenv()

client = openai.OpenAI(
    base_url=os.environ["OPENAI_API_BASE"],
    api_key=os.environ["OPENAI_API_KEY"],
)

# 超过这个数量的问题，需要场景引导
QUESTION_NUM_THRESHOLD = 10

DEFAULT_MODEL_NAME = os.getenv("MODEL_NAME")

PROMPT_TEMPLATE = """你是一个经验丰富的数据标注专家，现在需要你根据问题，提取出场景。

一级事项：{level_1}

二级事项：{level_2}

问题列表：
{question_list}

这部分类别的问题较多，我做RAG的时候，无法定位到具体的场景，需要通过反问来向用户澄清。因此我需要你帮我将这些问题进行归类。
输出格式为json，格式如下：
{{
    "场景引导类型1": [
        "场景引导标签1",
        "场景引导标签2",
        ...
    ],
    "场景引导类型2": [
        "场景引导标签1",
        "场景引导标签2",
        ...
    ],
    ...
}}
"""

EXTRACT_SCENE_PROMPT_TEMPLATE = """任务描述：
你是一个数据标注员，现在需要你根据问题，提取出场景。

候选场景：
{candidate_scene}

待抽取问题：
{question}

请根据候选场景，提取出问题对应的场景标签。其中候选场景中key是category，value是label。

输出格式：

```json
[
    {{
        "category": "场景引导类型1",
        "label": "场景引导标签1"
    }}
]
```
注意如果没有命中任何候选场景，请返回空列表。如果命中多个候选场景，请返回多个场景。
"""


@click.command()
@click.option("--input-file", "-i", type=str, help="输入文件")
@click.option("--env-file", "-e", type=str, help="大模型环境变量文件")
def main(input_file: str, env_file: str):
    current_dir = os.path.dirname(os.path.abspath(__file__))

    df = pd.read_excel(input_file)

    # group by level_1 and level_2
    groups = df.groupby(["level_1", "level_2"])

    for pair, group in groups:
        if len(group) < QUESTION_NUM_THRESHOLD:
            continue

        prompt = PROMPT_TEMPLATE.format(
            level_1=pair[0],
            level_2=pair[1],
            question_list="\n- ".join(group["question"].tolist()),
        )
        client = get_by_name(env_file)

        response = client(user_input=prompt)
        scenes = response.parse_json()
        print(scenes)

        for _, row in tqdm(group.iterrows(), total=len(group)):
            question = row["question"]
            extract_scene_prompt = EXTRACT_SCENE_PROMPT_TEMPLATE.format(
                level_1=pair[0],
                level_2=pair[1],
                candidate_scene=json.dumps(scenes, ensure_ascii=False, indent=2),
                question=question,
            )

            response = client(user_input=extract_scene_prompt)
            print(response.content)
            scenes = response.parse_json(is_list=True)

            scene_strs = []
            for scene in scenes:
                scene_obj = Scene(**scene)
                scene_str = scene_obj.to_str()
                scene_strs.append(scene_str)

            df.loc[row.name, "scene"] = "、".join(scene_strs)
            print(df.loc[row.name, "scene"])

    df.to_excel(input_file, index=False)


if __name__ == "__main__":
    main()